from matplotlib import pyplot as plt
from matplotlib.colors import ListedColormap
# from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
import time
from data.neu_seg_competition_nation.dataset import Neu_Seg_Competition_Dataset
from torch.utils.data import DataLoader
from utils import *


# writer = SummaryWriter("./tblogs")


def record_images_and_masks(writer, images, masks, pred, global_step, prefix="train"):
    """
    记录原始图像、真实掩膜和预测掩膜到 TensorBoard。

    参数:
    - writer (SummaryWriter): TensorBoard 的 SummaryWriter 对象。
    - images (torch.Tensor): 形状为 (batch_size, channels, height, width) 的原始图像张量。
    - masks (torch.Tensor): 形状为 (batch_size, height, width) 的真实掩膜张量。
    - pred_masks (torch.Tensor): 形状为 (batch_size, height, width) 的预测掩膜张量。
    - global_step (int): 当前全局步数。
    - prefix (str): 记录的前缀，默认为 "train"。
    """
    pred_masks = convert_pred_to_mask(pred)
    pred_masks_numpy = (pred_masks[0]).squeeze(0).cpu().numpy()
    masks_numpy = (masks[0]).cpu().numpy()
    image_numpy = (images[0]).permute(1, 2, 0).cpu().numpy()
    # # 将图像和掩膜组合成网格
    # grid_images = torchvision.utils.make_grid(images, nrow=4, padding=2, normalize=True, scale_each=True)
    # grid_masks = torchvision.utils.make_grid(masks, nrow=4, padding=2, normalize=True,
    #                                          scale_each=True)
    # grid_pred_masks = torchvision.utils.make_grid(pred_masks, nrow=4, padding=2,
    #                                               normalize=True, scale_each=True)

    # # 记录原始图像
    # writer.add_image(f"{prefix}/images", grid_images, global_step)
    #
    # # 记录真实掩膜
    # writer.add_image(f"{prefix}/masks", grid_masks, global_step)
    #
    # # 记录预测掩膜
    # writer.add_image(f"{prefix}/pred_masks", grid_pred_masks, global_step)

    # grid_images_np = grid_images.permute(1, 2, 0).numpy()
    # grid_masks_np = grid_masks.permute(1, 2, 0).numpy()
    # grid_pred_masks_np = grid_pred_masks.permute(1, 2, 0).numpy()
    # total_iou = {i: 0.0 for i in range(4)}
    # for i in range(4):
    #     ki,u = get_iandu_array(pred_masks_numpy, masks_numpy, i)
    #     if u == 0:
    #         continue
    #     total_iou[i] += float(ki)/float(u)
    # print(f'o: 1:{total_iou[1]}, 2:{total_iou[2]},3:{total_iou[3]},miou:{(total_iou[1]+total_iou[2]+total_iou[3])/3}')
    #
    # n2 , u2 = get_scores(pred[0].unsqueeze(0), masks[0].unsqueeze(0),device=device)
    # iouclass = n2/u2
    # print(f'me: 1:{iouclass[1].item()}, 2:{iouclass[2].item()},3:{iouclass[3].item()},miou:{(iouclass[1].item()+iouclass[2].item()+iouclass[3].item())/3}')
    return image_numpy, masks_numpy, pred_masks_numpy

def evaluate(name,pathi):
    batch_size = 1  # 评估模式，不能动这个
    num_classes = 4
    pin_memory = False
    num_workers = min(torch.cuda.device_count()*4,batch_size)
    # 检查可用的 GPU 数量
    device_ids = list(range(torch.cuda.device_count()))

    # 创建设备对象
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    print(f"Available GPUs: {device_ids}")

    # 初始化模型 miou_0.6526827664614653_model_best_epoch_4095.pth
    model = torch.load(pathi)
    # 根据 GPU 数量决定是否使用 DataParallel
    if len(device_ids) > 1:
        batch_size *= len(device_ids)
        model = torch.nn.DataParallel(model, device_ids=device_ids)

    # 将模型移动到指定设备
    model.to(device)

    train_dataset = Neu_Seg_Competition_Dataset(data_path='./data/neu_seg_competition', data_type='training')
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory)
    test_dataset = Neu_Seg_Competition_Dataset(data_path='./data/neu_seg_competition', data_type='test')
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers,
                             pin_memory=pin_memory)
    model.eval()

    all_sets = train_dataset + test_dataset
    all_loader = DataLoader(all_sets, batch_size=batch_size, shuffle=False, num_workers=num_workers,
                             pin_memory=pin_memory)
    # 初始化统计变量
    total_n = torch.zeros(num_classes, device=device)
    total_u = torch.zeros(num_classes, device=device)
    total_n1 = torch.zeros(num_classes, device=device)
    total_u1 = torch.zeros(num_classes, device=device)
    colors = ['black', 'red', 'green', 'blue']  # 0类，1类，2类，3类
    cmap = ListedColormap(colors)
    with torch.no_grad():
        idx = 0
        t_all = []
        for data in tqdm(all_loader):
            image, mask, image_name = data
            # 将数据移动到合适的设备（例如 GPU）
            image = image.to(device)
            mask = mask.to(device)
            t1 = time.time()
            # 前向传播
            outputs = model(image)
            t2 = time.time()
            t_all.append(t2 - t1)
            # 计算 IoU 和 mIoU
            ni, ui = get_scores(outputs, mask,device)
            # 累加 IoU
            ni = ni.to(device)
            ui = ui.to(device)
            uinew1 = torch.where(ui < 1, 1, ui)
            iou = ni / uinew1
            miou = (iou[1].item() + iou[2].item() + iou[3].item()) / 3
            total_n += ni
            total_u += ui
            idx+=1
            # print(image_name[0])
            # if image_name[0]!='In_165':
            #     continue
            # a, b, c = record_images_and_masks(111, image, mask, outputs, global_step=idx, prefix="test")
            # idx += 1
            # fig, axs = plt.subplots(2, 2, figsize=(15, 5))
            # # 绘制原始图像
            # axs[0,0].imshow(a)
            # axs[0,0].set_title(f'{image_name}')
            # axs[0,0].axis('off')

            # # 绘制真实掩膜
            # axs[0,1].imshow(b)
            # axs[0,1].set_title(f'True Masks{np.unique(b)}')
            # axs[0,1].axis('off')

            # # 绘制预测掩膜
            # axs[1,1].imshow(c)
            # axs[1,1].set_title(f'Predicted Masks/{miou}/{np.unique(c)}')
            # axs[1,1].axis('off')

            # #绘制 预测掩膜和真实掩膜差异
            # axs[1,0].imshow(np.abs(c - b) )
            # axs[1,0].set_title(f'abs(me-mask)/{np.unique(np.abs(c - b))}')
            # axs[1,0].axis('off')


            # # 显示图像
            # plt.tight_layout()
            # plt.show()

    # 计算平均 mIoU
    total_u = torch.where(total_u < 1, 1, total_u)
    total_n = torch.where(total_n < 1,1,total_n)
    class_iou = total_n / total_u
    avg_class1_iou = class_iou[1].item()
    avg_class2_iou = class_iou[2].item()
    avg_class3_iou = class_iou[3].item()
    avg_miou = (avg_class1_iou + avg_class2_iou + avg_class3_iou) / 3
    # print('average time:', np.mean(t_all) / 1)
    # print('average fps:', 1 / np.mean(t_all))
    fps = 1 / np.mean(t_all)
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print('[{}]avg_miou:{},avg_class1_iou{},avg_class2_iou{},avg_class3_iou{},参数量{},FPS{}'.format(name,avg_miou, avg_class1_iou,
                                                                                        avg_class2_iou, avg_class3_iou,total_params,fps))

if __name__ == '__main__':
    from model.UNet import *
    from model.E_net import *
    # evaluate('Unet','./save_model_best/Unet_baseline_miou_0.7565735578536987_model_best_epoch_34 (1).pth')
    evaluate('E_net','./save_model_best/ENet_enhance_b512_random_model_epoch_19999_loss_0.30629223585128784.pth')